{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/aburkov/theLMbook/blob/main/emotion_classifier_LR.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "<div style=\"display: flex; justify-content: center;\">\n",
        "    <div style=\"background-color: #f4f6f7; padding: 15px; width: 80%;\">\n",
        "        <table style=\"width: 100%\">\n",
        "            <tr>\n",
        "                <td style=\"vertical-align: middle;\">\n",
        "                    <span style=\"font-size: 14px;\">\n",
        "                        A notebook for <a href=\"https://www.thelmbook.com\" target=\"_blank\" rel=\"noopener\">The Hundred-Page Language Models Book</a> by Andriy Burkov<br><br>\n",
        "                        Code repository: <a href=\"https://github.com/aburkov/theLMbook\" target=\"_blank\" rel=\"noopener\">https://github.com/aburkov/theLMbook</a>\n",
        "                    </span>\n",
        "                </td>\n",
        "                <td style=\"vertical-align: middle;\">\n",
        "                    <a href=\"https://www.thelmbook.com\" target=\"_blank\" rel=\"noopener\">\n",
        "                        <img src=\"https://thelmbook.com/img/book.png\" width=\"80px\" alt=\"The Hundred-Page Language Models Book\">\n",
        "                    </a>\n",
        "                </td>\n",
        "            </tr>\n",
        "        </table>\n",
        "    </div>\n",
        "</div>"
      ],
      "metadata": {
        "id": "P9WVD-1ZAmYf"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "yy0zjL_2ouOU",
        "outputId": "63da21d9-c17e-42fc-d9fd-aa1041b38dba"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Number of training examples: 18000\n",
            "Number of test examples: 2000\n",
            "\n",
            "Train accuracy: 0.9854\n",
            "Test accuracy: 0.8855\n",
            "\n",
            "--- Better hyperparameters ---\n",
            "Train accuracy: 0.9962\n",
            "Test accuracy: 0.8910\n"
          ]
        }
      ],
      "source": [
        "# Import required libraries\n",
        "import gzip             # For decompressing gzipped data files\n",
        "import json             # For parsing JSON-formatted data\n",
        "import random           # For shuffling dataset and setting seeds\n",
        "import requests         # For downloading dataset from URL\n",
        "from sklearn.feature_extraction.text import CountVectorizer # Text vectorization utility\n",
        "from sklearn.linear_model import LogisticRegression         # Logistic regression model\n",
        "from sklearn.metrics import accuracy_score                  # For model evaluation\n",
        "\n",
        "# ----------------------------\n",
        "# Utility Functions\n",
        "# ----------------------------\n",
        "\n",
        "def set_seed(seed):\n",
        "    \"\"\"\n",
        "    Sets random seed for reproducibility.\n",
        "\n",
        "    Args:\n",
        "        seed (int): Seed value for random number generation\n",
        "    \"\"\"\n",
        "    random.seed(seed)\n",
        "\n",
        "def download_and_split_data(data_url, test_ratio=0.1):\n",
        "    \"\"\"\n",
        "    Downloads emotion classification dataset from URL and splits into train/test sets.\n",
        "    Handles decompression and JSON parsing of the raw data.\n",
        "\n",
        "    Args:\n",
        "        data_url (str): URL of the gzipped JSON dataset\n",
        "        test_ratio (float): Proportion of data to use for testing (default: 0.1)\n",
        "\n",
        "    Returns:\n",
        "        tuple: (X_train, y_train, X_test, y_test) containing:\n",
        "            - X_train, X_test: Lists of text examples for training and testing\n",
        "            - y_train, y_test: Lists of corresponding emotion labels\n",
        "    \"\"\"\n",
        "    # Download and decompress the dataset\n",
        "    response = requests.get(data_url)\n",
        "    content = gzip.decompress(response.content).decode()\n",
        "\n",
        "    # Parse JSON lines into list of dictionaries\n",
        "    dataset = [json.loads(line) for line in content.splitlines()]\n",
        "\n",
        "    # Shuffle dataset for random split\n",
        "    random.shuffle(dataset)\n",
        "\n",
        "    # Split into train and test sets\n",
        "    split_index = int(len(dataset) * (1 - test_ratio))\n",
        "    train, test = dataset[:split_index], dataset[split_index:]\n",
        "\n",
        "    # Separate text and labels\n",
        "    X_train = [item[\"text\"] for item in train]\n",
        "    y_train = [item[\"label\"] for item in train]\n",
        "    X_test = [item[\"text\"] for item in test]\n",
        "    y_test = [item[\"label\"] for item in test]\n",
        "\n",
        "    return X_train, y_train, X_test, y_test\n",
        "\n",
        "# ----------------------------\n",
        "# Main Execution\n",
        "# ----------------------------\n",
        "\n",
        "# Set random seed for reproducibility\n",
        "set_seed(42)\n",
        "\n",
        "# Download and prepare dataset\n",
        "data_url = \"https://www.thelmbook.com/data/emotions\"\n",
        "X_train_text, y_train, X_test_text, y_test = download_and_split_data(\n",
        "    data_url, test_ratio=0.1\n",
        ")\n",
        "\n",
        "print(\"Number of training examples:\", len(X_train_text))\n",
        "print(\"Number of test examples:\", len(X_test_text))\n",
        "\n",
        "# ----------------------------\n",
        "# Baseline Model\n",
        "# ----------------------------\n",
        "\n",
        "# Initialize text vectorizer with basic parameters\n",
        "# max_features=10_000: Limit vocabulary to top 10k most frequent words\n",
        "# binary=True: Convert counts to binary indicators (0/1)\n",
        "vectorizer = CountVectorizer(max_features=10_000, binary=True)\n",
        "\n",
        "# Transform text data to numerical features\n",
        "X_train = vectorizer.fit_transform(X_train_text)\n",
        "X_test = vectorizer.transform(X_test_text)\n",
        "\n",
        "# Initialize and train logistic regression model\n",
        "model = LogisticRegression(random_state=42, max_iter=1000)\n",
        "model.fit(X_train, y_train)\n",
        "\n",
        "# Make predictions on train and test sets\n",
        "y_train_pred = model.predict(X_train)\n",
        "y_test_pred = model.predict(X_test)\n",
        "\n",
        "# Calculate and display accuracy metrics\n",
        "train_accuracy = accuracy_score(y_train, y_train_pred)\n",
        "test_accuracy = accuracy_score(y_test, y_test_pred)\n",
        "\n",
        "print(f\"\\nTrain accuracy: {train_accuracy:.4f}\")\n",
        "print(f\"Test accuracy: {test_accuracy:.4f}\")\n",
        "\n",
        "# ----------------------------\n",
        "# Improved Model\n",
        "# ----------------------------\n",
        "\n",
        "print(\"\\n--- Better hyperparameters ---\")\n",
        "\n",
        "# Initialize vectorizer with improved parameters\n",
        "# max_features=20000: Increased vocabulary size\n",
        "# ngram_range=(1, 2): Include both unigrams and bigrams\n",
        "vectorizer = CountVectorizer(max_features=20000, ngram_range=(1, 2))\n",
        "\n",
        "# Transform text data with new vectorizer\n",
        "X_train = vectorizer.fit_transform(X_train_text)\n",
        "X_test = vectorizer.transform(X_test_text)\n",
        "\n",
        "# Train and evaluate model with same parameters\n",
        "model = LogisticRegression(random_state=42, max_iter=1000)\n",
        "model.fit(X_train, y_train)\n",
        "\n",
        "y_train_pred = model.predict(X_train)\n",
        "y_test_pred = model.predict(X_test)\n",
        "\n",
        "train_accuracy = accuracy_score(y_train, y_train_pred)\n",
        "test_accuracy = accuracy_score(y_test, y_test_pred)\n",
        "\n",
        "print(f\"Train accuracy: {train_accuracy:.4f}\")\n",
        "print(f\"Test accuracy: {test_accuracy:.4f}\")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyNKH13VydD5aNyFkNnbyP3F",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}